/*******************************************************
 * Copyright (c) 2015, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <arrayfire.h>
#include <cmath>
#include <cstdio>
#include <cstdlib>

using namespace af;

/**
 * randomization - controls % of total number of pixels in the image
 * that will be effected by random noise
 * repeat - # of times the process is carried out on the previous steps output
 */
array hurl(const array &in, int randomization, int repeat) {
    int w         = in.dims(0);
    int h         = in.dims(1);
    float f       = randomization / 100.0f;
    int dim       = (int)(f * w * h);
    array ret_val = in.copy();
    array temp    = moddims(ret_val, w * h, 3);
    for (int i = 0; i < repeat; ++i) {
        array idxs    = (w * h) * randu(dim);
        array rndR    = 255.0f * randu(dim);
        array rndG    = 255.0f * randu(dim);
        array rndB    = 255.0f * randu(dim);
        temp(idxs, 0) = rndR;
        temp(idxs, 1) = rndG;
        temp(idxs, 2) = rndB;
    }
    ret_val = moddims(temp, in.dims());
    return ret_val;
}

/**
 * Retrieve a new image of same dimensions of the original image
 * where each original image's pixel is replaced by randomly picked
 * neighbor in the provided local neighborhood window
 */
array getRandomNeighbor(const array &in, int windW, int windH) {
    array rnd = 2.0f * randu(in.dims(0), in.dims(1)) - 1.0f;
    array sx  = seq(in.dims(0));
    array sy  = seq(in.dims(1));
    array vx  = tile(sx, 1, in.dims(1)) + floor(rnd * windW);
    array vy  = tile(sy.T(), in.dims(0), 1) + floor(rnd * windH);
    array vxx = clamp(vx, 0, in.dims(0));
    array vyy = clamp(vy, 0, in.dims(1));
    array in2 = moddims(in, vx.elements(), 3);
    return moddims(in2(vyy * in.dims(0) + vxx, span), in.dims());
}

/**
 * randomly pick neighbor from given window size and replace the
 * current pixel with the randomly chosen color.
 * No new colors are introduced, unlike hurl.
 */
array spread(const array &in, int window_width, int window_height) {
    return getRandomNeighbor(in, window_width, window_height);
}

/**
 * randomization - controls % of total number of pixels in the image
 * that will be effected by random noise
 * repeat - # of times the process is carried out on the previous steps output
 */
array pick(const array &in, int randomization, int repeat) {
    int w         = in.dims(0);
    int h         = in.dims(1);
    float f       = randomization / 100.0f;
    int dim       = (int)(f * w * h);
    array ret_val = in.copy();
    for (int i = 0; i < repeat; ++i) {
        array idxs           = (w * h) * randu(dim);
        array rnd            = getRandomNeighbor(ret_val, 1, 1);
        array temp_src       = moddims(rnd, w * h, 3);
        array temp_dst       = moddims(ret_val, w * h, 3);
        temp_dst(idxs, span) = temp_src(idxs, span);
        ret_val              = moddims(temp_dst, in.dims());
    }
    return ret_val;
}

void prewitt(array &mag, array &dir, const array &in) {
    static float h1[] = {1, 1, 1};
    static float h2[] = {-1, 0, 1};
    static array h1d(3, h1);
    static array h2d(3, h2);

    // Find the gradients
    array Gy = af::convolve(h2d, h1d, in) / 6;
    array Gx = af::convolve(h1d, h2d, in) / 6;

    // Find magnitude and direction
    mag = hypot(Gx, Gy);
    dir = atan2(Gy, Gx);
}

void sobelFilter(array &mag, array &dir, const array &in) {
    // Find the gradients
    array Gy, Gx;
    af::sobel(Gx, Gy, in);

    // Find magnitude and direction
    mag = hypot(Gx, Gy);
    dir = atan2(Gy, Gx);
}

void normalizeImage(array &in) {
    float min = af::min<float>(in);
    float max = af::max<float>(in);
    in        = 255.0f * ((in - min) / (max - min));
}

array DifferenceOfGaussian(const array &in, int window_radius1,
                           int window_radius2) {
    array ret_val;
    int w1   = 2 * window_radius1 + 1;
    int w2   = 2 * window_radius2 + 1;
    array g1 = gaussianKernel(w1, w1);
    array g2 = gaussianKernel(w2, w2);
    ret_val  = (convolve(in, g1) - convolve(in, g2));
    normalizeImage(ret_val);
    return ret_val;
}

array medianfilter(const array &in, int window_width, int window_height) {
    array ret_val(in.dims());
    ret_val(span, span, 0) =
        medfilt(in(span, span, 0), window_width, window_height);
    ret_val(span, span, 1) =
        medfilt(in(span, span, 1), window_width, window_height);
    ret_val(span, span, 2) =
        medfilt(in(span, span, 2), window_width, window_height);
    return ret_val;
}

array gaussianblur(const array &in, int window_width, int window_height,
                   double sigma) {
    array g = gaussianKernel(window_width, window_height, sigma, sigma);
    return convolve(in, g);
}

/**
 * azimuth range is [0-360]
 * elevation range is [0-180]
 * depth range is [1-100]
 * Note: this function has been tailored after
 * the emboss implementation in GIMP editor
 **/
array emboss(const array &input, float azimuth, float elevation, float depth) {
    if (depth < 1 || depth > 100) {
        printf("Depth should be in the range of 1-100");
        return input;
    }
    static float x[3] = {-1, 0, 1};
    static array hg(3, x);
    static array vg = hg.T();

    array in = input;
    if (in.dims(2) > 1)
        in = colorSpace(input, AF_GRAY, AF_RGB);
    else
        in = input;

    // convert angles to radians
    float phi   = elevation * af::Pi / 180.0f;
    float theta = azimuth * af::Pi / 180.0f;

    // compute light pos in cartesian coordinates
    // and scale with maximum intensity
    // phi will effect the amount of we intend to put
    // on a pixel
    float pos[3];
    pos[0] = 255.99f * cos(phi) * cos(theta);
    pos[1] = 255.99f * cos(phi) * sin(theta);
    pos[2] = 255.99f * sin(phi);

    // compute gradient vector
    array gx = convolve(in, vg);
    array gy = convolve(in, hg);

    float pxlz   = (6 * 255.0f) / depth;
    array zdepth = constant(pxlz, gx.dims());
    array vdot   = gx * pos[0] + gy * pos[1] + pxlz * pos[2];
    array outwd  = vdot < 0.0f;
    array norm   = vdot / sqrt(gx * gx + gy * gy + zdepth * zdepth);

    array color = outwd * 0.0f + (1 - outwd) * norm;
    return color;
}

int main(int argc, char **argv) {
    try {
        int device = argc > 1 ? atoi(argv[1]) : 0;
        af::setDevice(device);
        af::info();

        array img =
            loadImage(ASSETS_DIR "/examples/images/vegetable-woman.jpg", true);

        array prew_mag, prew_dir;
        array sob_mag, sob_dir;
        array img1ch = colorSpace(img, AF_GRAY, AF_RGB);
        prewitt(prew_mag, prew_dir, img1ch);
        sobelFilter(sob_mag, sob_dir, img1ch);
        array sprd  = spread(img, 3, 3);
        array hrl   = hurl(img, 10, 1);
        array pckng = pick(img, 40, 2);
        array difog = DifferenceOfGaussian(img, 1, 2);
        array bil   = bilateral(hrl, 3.0f, 40.0f);
        array mf    = medianfilter(hrl, 5, 5);
        array gb    = gaussianblur(hrl, 3, 3, 0.8);
        array emb   = emboss(img, 45, 20, 10);

        af::Window wnd("Image Filters Demo");
        printf("Press ESC while the window is in focus to exit\n");
        while (!wnd.close()) {
            wnd.grid(2, 5);
            wnd(0, 0).image(hrl / 255, "Hurl noise");
            wnd(1, 0).image(gb / 255, "Gaussian blur");
            wnd(0, 1).image(bil / 255, "Bilateral filter on hurl noise");
            wnd(1, 1).image(mf / 255, "Median filter on hurl noise");
            wnd(0, 2).image(prew_mag / 255, "Prewitt edge filter");
            wnd(1, 2).image(sob_mag / 255, "Sobel edge filter");
            wnd(0, 3).image(sprd / 255, "Spread filter");
            wnd(1, 3).image(pckng / 255, "Pick filter");
            wnd(0, 4).image(difog / 255,
                            "Difference of gaussians(3x3 and 5x5)");
            wnd(1, 4).image(emb / 255, "Emboss effect");
            wnd.show();
        }

    } catch (af::exception &e) {
        fprintf(stderr, "%s\n", e.what());
        throw;
    }
    return 0;
}
